
from paddleocr import PaddleOCR

import numpy as np
import cv2

# paddle_ocr.py中配置
class PaddleOCRProcessor:
    # def __init__(self):
    #     self.ocr = PaddleOCR(
    #         use_angle_cls=True,
    #         lang='ch',
    #         use_gpu=True,  # GPU加速
    #         rec_algorithm='SVTR_LCNet',
    #         det_db_score_mode='slow',
    #         max_text_length=100
    #     )


    def __init__(self, use_gpu=False):
        # 指定文件夹位置
        FILE_PATH = "E:/pub/python/python-ocr-demo"
        self.ocr = PaddleOCR(
            use_angle_cls=True,
            lang="ch",
            use_gpu=use_gpu,
            det_model_dir= f"{FILE_PATH}/models/trained_models/det/",
            rec_model_dir= f"{FILE_PATH}/models/trained_models/rec/",
        )

    # def __init__(self, use_gpu=False):
    #     self.ocr = PaddleOCR(use_angle_cls=True, lang="ch", use_gpu=use_gpu)

    def process_image(self, image_path):
        """核心处理方法"""
        # 读取图片
        img = cv2.imread(image_path)

        # 执行OCR
        result = self.ocr.ocr(img, cls=True)

        # 格式化输出
        formatted_results = []
        if result[0]:
            for line in result[0]:
                text = line[1][0]
                confidence = round(line[1][1], 4)
                coordinates = [list(map(int, point)) for point in line[0]]
                formatted_results.append({
                    "text": text,
                    "confidence": confidence,
                    "position": coordinates
                })
        return formatted_results


    def process_image1(self, file):
        """核心处理方法"""
        # 读取图片
        img = cv2.imdecode(np.frombuffer(file.read(), np.uint8), cv2.IMREAD_COLOR)
        # 执行OCR
        result = self.ocr.ocr(img, cls=True)

        # 格式化输出
        formatted_results = []
        if result[0]:
            for line in result[0]:
                text = line[1][0]
                confidence = round(line[1][1], 4)
                coordinates = [list(map(int, point)) for point in line[0]]
                formatted_results.append({
                    "text": text,
                    "confidence": confidence,
                    "position": coordinates
                })
        return formatted_results

# 方式2：使用挂载的自定义路径
# ocr = PaddleOCR(
#     det_model_dir=os.getenv('MODEL_DIR')+'/inference/det',
#     rec_model_dir=os.getenv('MODEL_DIR')+'/inference/rec'
# )


